// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

// Useful operations to implement kernels handling run-end encoded arrays

#include <algorithm>
#include <cstdint>
#include <limits>
#include <memory>

#include <nebula/array/data.h>
#include <nebula/compute/exec.h>
#include <nebula/compute/kernel.h>

#include <turbo/utility/status.h>
#include <nebula/types/type_traits.h>
#include <nebula/bits/bit_util.h>
#include <turbo/log/logging.h>

namespace nebula::compute::internal::ree_util {

    template<typename ArrowType, bool in_has_validity_buffer,
            bool out_has_validity_buffer = in_has_validity_buffer, typename Enable = void>
    struct ReadWriteValue {
    };

    // Numeric and primitive C-compatible types
    template<typename ArrowType, bool in_has_validity_buffer, bool out_has_validity_buffer>
    class ReadWriteValue<ArrowType, in_has_validity_buffer, out_has_validity_buffer,
            enable_if_has_c_type<ArrowType>> {
    public:
        using ValueRepr = typename ArrowType::c_type;

    private:
        const uint8_t *input_validity_;
        const uint8_t *input_values_;

        // Needed only by the writing functions
        uint8_t *output_validity_;
        uint8_t *output_values_;

    public:
        explicit ReadWriteValue(const ArraySpan &input_values_array,
                                ArrayData *output_values_array_data)
                : input_validity_(in_has_validity_buffer ? input_values_array.buffers[0].data
                                                         : nullptr),
                  input_values_(input_values_array.buffers[1].data),
                  output_validity_((out_has_validity_buffer && output_values_array_data)
                                   ? output_values_array_data->buffers[0]->mutable_data()
                                   : nullptr),
                  output_values_(output_values_array_data
                                 ? output_values_array_data->buffers[1]->mutable_data()
                                 : nullptr) {}

        [[nodiscard]] bool ReadValue(ValueRepr *out, int64_t read_offset) const {
            bool valid = true;
            if constexpr (in_has_validity_buffer) {
                valid = bit_util::get_bit(input_validity_, read_offset);
            }
            if constexpr (std::is_same_v<ArrowType, BooleanType>) {
                *out = bit_util::get_bit(input_values_, read_offset);
            } else {
                *out = (reinterpret_cast<const ValueRepr *>(input_values_))[read_offset];
            }
            return valid;
        }

        /// Pre-conditions guaranteed by the callers:
        /// - i and j are valid indices into the values buffer
        /// - the values in i and j are valid
        bool CompareValuesAt(int64_t i, int64_t j) const {
            if constexpr (std::is_same_v<ArrowType, BooleanType>) {
                return bit_util::get_bit(input_values_, i) == bit_util::get_bit(input_values_, j);
            } else {
                return (reinterpret_cast<const ValueRepr *>(input_values_))[i] ==
                       (reinterpret_cast<const ValueRepr *>(input_values_))[j];
            }
        }

        /// \brief Ensure padding is zeroed in validity bitmap.
        void ZeroValidityPadding(int64_t length) const {
            DKCHECK(output_values_);
            if constexpr (out_has_validity_buffer) {
                DKCHECK(output_validity_);
                const int64_t validity_buffer_size = bit_util::BytesForBits(length);
                output_validity_[validity_buffer_size - 1] = 0;
            }
        }

        void WriteValue(int64_t write_offset, bool valid, ValueRepr value) const {
            if constexpr (out_has_validity_buffer) {
                bit_util::SetBitTo(output_validity_, write_offset, valid);
            }
            if (valid) {
                if constexpr (std::is_same_v<ArrowType, BooleanType>) {
                    bit_util::SetBitTo(output_values_, write_offset, value);
                } else {
                    (reinterpret_cast<ValueRepr *>(output_values_))[write_offset] = value;
                }
            }
        }

        void WriteRun(int64_t write_offset, int64_t run_length, bool valid,
                      ValueRepr value) const {
            if constexpr (out_has_validity_buffer) {
                bit_util::SetBitsTo(output_validity_, write_offset, run_length, valid);
            }
            if (valid) {
                if constexpr (std::is_same_v<ArrowType, BooleanType>) {
                    bit_util::SetBitsTo(reinterpret_cast<uint8_t *>(output_values_), write_offset,
                                        run_length, value);
                } else {
                    auto *output_values_c = reinterpret_cast<ValueRepr *>(output_values_);
                    std::fill(output_values_c + write_offset,
                              output_values_c + write_offset + run_length, value);
                }
            }
        }

        bool compare(ValueRepr lhs, ValueRepr rhs) const { return lhs == rhs; }
    };

    // FixedSizeBinary, Decimal128
    template<typename ArrowType, bool in_has_validity_buffer, bool out_has_validity_buffer>
    class ReadWriteValue<ArrowType, in_has_validity_buffer, out_has_validity_buffer,
            enable_if_fixed_size_binary<ArrowType>> {
    public:
        // Every value is represented as a pointer to byte_width_ bytes
        using ValueRepr = uint8_t const *;

    private:
        const uint8_t *input_validity_;
        const uint8_t *input_values_;

        // Needed only by the writing functions
        uint8_t *output_validity_;
        uint8_t *output_values_;

        const size_t byte_width_;

    public:
        ReadWriteValue(const ArraySpan &input_values_array, ArrayData *output_values_array_data)
                : input_validity_(in_has_validity_buffer ? input_values_array.buffers[0].data
                                                         : nullptr),
                  input_values_(input_values_array.buffers[1].data),
                  output_validity_((out_has_validity_buffer && output_values_array_data)
                                   ? output_values_array_data->buffers[0]->mutable_data()
                                   : nullptr),
                  output_values_(output_values_array_data
                                 ? output_values_array_data->buffers[1]->mutable_data()
                                 : nullptr),
                  byte_width_(input_values_array.type->byte_width()) {}

        [[nodiscard]] bool ReadValue(ValueRepr *out, int64_t read_offset) const {
            bool valid = true;
            if constexpr (in_has_validity_buffer) {
                valid = bit_util::get_bit(input_validity_, read_offset);
            }
            *out = input_values_ + (read_offset * byte_width_);
            return valid;
        }

        bool CompareValuesAt(int64_t i, int64_t j) const {
            return 0 == memcmp(input_values_ + (i * byte_width_),
                               input_values_ + (j * byte_width_), byte_width_);
        }

        /// \brief Ensure padding is zeroed in validity bitmap.
        void ZeroValidityPadding(int64_t length) const {
            DKCHECK(output_values_);
            if constexpr (out_has_validity_buffer) {
                DKCHECK(output_validity_);
                const int64_t validity_buffer_size = bit_util::BytesForBits(length);
                output_validity_[validity_buffer_size - 1] = 0;
            }
        }

        void WriteValue(int64_t write_offset, bool valid, ValueRepr value) const {
            if constexpr (out_has_validity_buffer) {
                bit_util::SetBitTo(output_validity_, write_offset, valid);
            }
            if (valid) {
                memcpy(output_values_ + (write_offset * byte_width_), value, byte_width_);
            }
        }

        void WriteRun(int64_t write_offset, int64_t run_length, bool valid,
                      ValueRepr value) const {
            if constexpr (out_has_validity_buffer) {
                bit_util::SetBitsTo(output_validity_, write_offset, run_length, valid);
            }
            if (valid) {
                uint8_t *ptr = output_values_ + (write_offset * byte_width_);
                for (int64_t i = 0; i < run_length; ++i) {
                    memcpy(ptr, value, byte_width_);
                    ptr += byte_width_;
                }
            }
        }

        bool compare(ValueRepr lhs, ValueRepr rhs) const {
            return memcmp(lhs, rhs, byte_width_) == 0;
        }
    };

    // Binary, String...
    template<typename ArrowType, bool in_has_validity_buffer, bool out_has_validity_buffer>
    class ReadWriteValue<ArrowType, in_has_validity_buffer, out_has_validity_buffer,
            enable_if_base_binary<ArrowType>> {
    public:
        using ValueRepr = std::string_view;
        using offset_type = typename ArrowType::offset_type;

    private:
        const uint8_t *input_validity_;
        const offset_type *input_offsets_;
        const uint8_t *input_values_;

        // Needed only by the writing functions
        uint8_t *output_validity_;
        offset_type *output_offsets_;
        uint8_t *output_values_;

    public:
        ReadWriteValue(const ArraySpan &input_values_array, ArrayData *output_values_array_data)
                : input_validity_(in_has_validity_buffer ? input_values_array.buffers[0].data
                                                         : nullptr),
                  input_offsets_(input_values_array.template get_values<offset_type>(1, 0)),
                  input_values_(input_values_array.buffers[2].data),
                  output_validity_((out_has_validity_buffer && output_values_array_data)
                                   ? output_values_array_data->buffers[0]->mutable_data()
                                   : nullptr),
                  output_offsets_(
                          output_values_array_data
                          ? output_values_array_data->template GetMutableValues<offset_type>(1, 0)
                          : nullptr),
                  output_values_(output_values_array_data
                                 ? output_values_array_data->buffers[2]->mutable_data()
                                 : nullptr) {}

        [[nodiscard]] bool ReadValue(ValueRepr *out, int64_t read_offset) const {
            bool valid = true;
            if constexpr (in_has_validity_buffer) {
                valid = bit_util::get_bit(input_validity_, read_offset);
            }
            if (valid) {
                const offset_type offset0 = input_offsets_[read_offset];
                const offset_type offset1 = input_offsets_[read_offset + 1];
                *out = std::string_view(reinterpret_cast<const char *>(input_values_ + offset0),
                                        offset1 - offset0);
            }
            return valid;
        }

        bool CompareValuesAt(int64_t i, int64_t j) const {
            const offset_type len_i = input_offsets_[i + 1] - input_offsets_[i];
            const offset_type len_j = input_offsets_[j + 1] - input_offsets_[j];
            return len_i == len_j &&
                   memcmp(input_values_ + input_offsets_[i], input_values_ + input_offsets_[j],
                          static_cast<size_t>(len_i));
        }

        /// \brief Ensure padding is zeroed in validity bitmap.
        void ZeroValidityPadding(int64_t length) const {
            DKCHECK(output_values_);
            if constexpr (out_has_validity_buffer) {
                DKCHECK(output_validity_);
                const int64_t validity_buffer_size = bit_util::BytesForBits(length);
                output_validity_[validity_buffer_size - 1] = 0;
            }
        }

        void WriteValue(int64_t write_offset, bool valid, ValueRepr value) const {
            if constexpr (out_has_validity_buffer) {
                bit_util::SetBitTo(output_validity_, write_offset, valid);
            }
            const offset_type offset0 = output_offsets_[write_offset];
            const offset_type offset1 =
                    offset0 + (valid ? static_cast<offset_type>(value.size()) : 0);
            output_offsets_[write_offset + 1] = offset1;
            if (valid) {
                memcpy(output_values_ + offset0, value.data(), value.size());
            }
        }

        void WriteRun(int64_t write_offset, int64_t run_length, bool valid,
                      ValueRepr value) const {
            if constexpr (out_has_validity_buffer) {
                bit_util::SetBitsTo(output_validity_, write_offset, run_length, valid);
            }
            if (valid) {
                int64_t i = write_offset;
                offset_type offset = output_offsets_[i];
                while (i < write_offset + run_length) {
                    memcpy(output_values_ + offset, value.data(), value.size());
                    offset += static_cast<offset_type>(value.size());
                    i += 1;
                    output_offsets_[i] = offset;
                }
            } else {
                offset_type offset = output_offsets_[write_offset];
                offset_type *begin = output_offsets_ + write_offset + 1;
                std::fill(begin, begin + run_length, offset);
            }
        }

        bool compare(ValueRepr lhs, ValueRepr rhs) const { return lhs == rhs; }
    };

    turbo::Result<std::shared_ptr<Buffer>> AllocateValuesBuffer(int64_t length, const DataType &type,
                                                                MemoryPool *pool,
                                                                int64_t data_buffer_size);

    turbo::Result<std::shared_ptr<ArrayData>> PreallocateRunEndsArray(
            const std::shared_ptr<DataType> &run_end_type, int64_t physical_length,
            MemoryPool *pool);

    /// \brief Preallocate the physical values array for a run-end encoded array
    ///
    /// data_buffer_size is passed here pre-calculated so this function doesn't have
    /// to be template-specialized for each type.
    ///
    /// The null_count is left as kUnknownNullCount (or 0 if length is 0) and, if
    /// after writing the values, the caller knows the null count, it can be set.
    ///
    /// \post if has_validity_buffer and length > 0, then data.buffer[0] != nullptr
    ///
    /// \param has_validity_buffer a validity buffer must be allocated
    /// \param length the length of the values array
    /// \param data_buffer_size the size of the data buffer for string and binary types
    turbo::Result<std::shared_ptr<ArrayData>> PreallocateValuesArray(
            const std::shared_ptr<DataType> &value_type, bool has_validity_buffer, int64_t length,
            MemoryPool *pool, int64_t data_buffer_size);

    /// \brief Preallocate the ArrayData for the run-end encoded version
    /// of the flat input array
    ///
    /// The top-level null_count is set to 0 (REEs keep all the data in child
    /// arrays). The null_count of the values array (child_data[1]) is left as
    /// kUnknownNullCount (or 0 if physical_length is 0) and, if after writing
    /// the values, the caller knows the null count, it can be set.
    ///
    /// \post if has_validity_buffer and physical_length > 0, then
    /// data.child_data[1].buffer[0] != nullptr
    ///
    /// \param data_buffer_size the size of the data buffer for string and binary types
    turbo::Result<std::shared_ptr<ArrayData>> PreallocateREEArray(
            std::shared_ptr<RunEndEncodedType> ree_type, bool has_validity_buffer,
            int64_t logical_length, int64_t physical_length, MemoryPool *pool,
            int64_t data_buffer_size);

    /// \brief Writes a single run-end to the first slot of the pre-allocated
    /// run-end encoded array in out
    ///
    /// Pre-conditions:
    /// - run_ends_data is of a valid run-ends type
    /// - run_ends_data has at least one slot
    /// - run_end > 0
    /// - run_ends fits in the run-end type without overflow
    void WriteSingleRunEnd(ArrayData *run_ends_data, int64_t run_end);

    turbo::Result<std::shared_ptr<ArrayData>> MakeNullREEArray(
            const std::shared_ptr<DataType> &run_end_type, int64_t logical_length,
            MemoryPool *pool);

}  // namespace  nebula::compute::internal::ree_util
